from osgeo import ogr
from osgeo import gdal
import numpy as np
import os
import geopandas


def RasterzeTheVectorToRasterr(inputfilePath, outputfile, templetefile):
    # 参数说明： 输出的栅格数据，注意该数据必须以update模式打开、指定要更新的波段个数(更新123波段)、指定的图层、几何图形坐标转换图像行列号函数、几何图形坐标转换图像行列号参数、以及图层中属性字段属性
    data = gdal.Open(templetefile)
    x_res = data.RasterXSize
    y_res = data.RasterYSize
    vector = ogr.Open(inputfilePath)
    layer = vector.GetLayer()
    targetDataSet = gdal.GetDriverByName('GTiff').Create(outputfile, x_res, y_res, 1, gdal.GDT_Float64)
    # targetDataSet=gdal.GetDriverByName('GTiff').CreateCopy(templetefile,data)
    targetDataSet.SetGeoTransform(data.GetGeoTransform())
    targetDataSet.SetProjection(data.GetProjection())
    band = targetDataSet.GetRasterBand(1)
    NoData_value = 0
    band.SetNoDataValue(NoData_value)
    band.FlushCache()
    gdal.RasterizeLayer(targetDataSet, [1], layer, options=["ATTRIBUTE=num_total"])


file_dir = 'E:/时空数据分析/o_monthly_json'

all_file_list = os.listdir(file_dir)
for single_file in all_file_list:
    inputfilePath = 'E:/时空数据分析/o_monthly_json/' + str(single_file) + ''
    outputfile = 'E:/时空数据分析/o_monthly_raster/' + str(single_file[:-8]) + '.tif'
    templetefile = 'E:/时空数据分析/参考.tif'
    RasterzeTheVectorToRasterr(inputfilePath, outputfile, templetefile)
    print(single_file)
